import time

from proto.stream import stream_pb2, stream_pb2_grpc
import grpc
from concurrent import futures


class Greeter(stream_pb2_grpc.GreeterServicer):

    def GetStream(self, request, context):
        idx = 1
        while context.is_active():
            idx += 1
            # close long connect
            if idx == 10:
                context.cancel()
            data = request.data
            time.sleep(1)
            yield stream_pb2.StreamResData(data="send {}/{}".format(data, idx))


    def PostGetStream(self, request_iterator, context):
        idx = 1
        for item in request_iterator:
            idx += 1
            print(item)
            if idx == 10:
                break
        return stream_pb2.StreamResData(data="OK")

    def AllGetStream(self, request_iterator, context):
        idx = 1
        for item in request_iterator:
            idx += 1
            print(item)
            if idx == 10:
                break
            yield stream_pb2.StreamResData(data=item.data)


if __name__ == '__main__':
    # instance server
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    # register server
    stream_pb2_grpc.add_GreeterServicer_to_server(Greeter(), server)
    # start
    server.add_insecure_port('0.0.0.0:5001')
    print("start service user...")
    server.start()
    server.wait_for_termination()
